//
// Created by hjhu on 2024/7/2.
//

#include "HourGlass.h"
#include "ConvBlock.h"

namespace coastal {

    HourGlassImpl::HourGlassImpl(int64_t num_modules, int64_t depth, int64_t num_features) : num_modules(num_modules),
                                                                                     depth(depth),
                                                                                     features(num_features) {
        generate_network(depth);
    }


    void HourGlassImpl::generate_network(int64_t level) {
        register_module("b1_" + std::to_string(level), std::make_shared<coastal::ConvBlockImpl>(features, features));
        register_module("b2_" + std::to_string(level), std::make_shared<coastal::ConvBlockImpl>(features, features));

        if (level > 1) {
            generate_network(level - 1);
        } else {
            register_module("b2_plus_" + std::to_string(level),
                            std::make_shared<coastal::ConvBlockImpl>(features, features));
        }

        register_module("b3_" + std::to_string(level), std::make_shared<coastal::ConvBlockImpl>(features, features));
    }

    torch::Tensor HourGlassImpl::_forward(int64_t level, torch::Tensor inp) {
        // Upper branch
        auto up1 = inp;
        up1 = get_module<ConvBlockImpl>("b1_" + std::to_string(level))->forward(up1);

        // Lower branch
        auto low1 = torch::avg_pool2d(inp, 2, 2);
        low1 = get_module<ConvBlockImpl>("b2_" + std::to_string(level))->forward(low1);

        torch::Tensor low2;
        if (level > 1) {
            low2 = _forward(level - 1, low1);
        } else {
            low2 = low1;
            low2 = get_module<ConvBlockImpl>("b2_plus_" + std::to_string(level))->forward(low2);
        }

        auto low3 = low2;
        low3 = get_module<ConvBlockImpl>("b3_" + std::to_string(level))->forward(low3);

        auto up2 = torch::upsample_nearest2d(low3, {low3.size(2) * 2, low3.size(3) * 2});

        return up1 + up2;
    }

    torch::Tensor HourGlassImpl::forward(torch::Tensor x) {
        return _forward(depth, x);
    }


} // coastal